// -*- coding: utf-8 -*- 
/**
 * Project: AlgorithmsLearn
 *
 * @author: yanking
 * Create time: 2022-04-03 00:31
 * IDE: IntelliJ IDEA
 * Introduction:
 */
package com.DataStruct.Tree.RBTree;

/**
 * 红黑树类
 *
 * @param <K>:红黑树关键字
 * @param <V>:红黑树关键字对应值
 */
class RBTree<K extends Comparable<K>, V> {
    private static final boolean RED = false;
    private static final boolean BLACK = true;
    private RBNode<K, V> root;

    private void inorder(RBNode<K, V> root) {
        if (root != null) {
            inorder(root.left);
            System.out.print(root.key + " " + (root.color ? "B" : "R") + " ");
            inorder(root.right);
        }
    }

    private void preorder(RBNode<K, V> root) {
        if (root != null) {
            System.out.print(root.key + " " + (root.color ? "B" : "R") + " ");
            preorder(root.left);
            preorder(root.right);
        }
    }

    public void preorderTraverse() {
        preorder(this.root);
        System.out.println();
    }

    public void inorderTraverse() {
        inorder(this.root);
        System.out.println();
    }

    public boolean colorOf(RBNode<K, V> node) {
        return node != null ? node.color : BLACK;
    }

    public void insert(K key, V value) {
        RBNode<K, V> t = this.root;
        if (t == null) {
            this.root = new RBNode<K, V>(null, key, value);
            setColor(root, BLACK);
            return;
        }
        int cmp = 0;
        RBNode<K, V> parent = null;
        while (t != null) {
            parent = t;
            cmp = key.compareTo(t.key);
            if (cmp < 0) {
                t = t.left;
            } else if (cmp > 0) {
                t = t.right;
            } else {
                t.value = value;
                return;
            }
        }
        RBNode<K, V> e = new RBNode<K, V>(parent, key, value);
        if (cmp < 0) {
            parent.left = e;
        } else {
            parent.right = e;
        }
        // 平衡处理，旋转+变色
        fixAfterInsert(e);
    }

    private void fixAfterInsert(RBNode<K, V> x) {
        setColor(x, RED);
        while (x != null && x != root && colorOf(x.parent) == RED) {
            if (x.parent == x.parent.parent.left) {
                RBNode<K, V> uncle = x.parent.parent.right;
                if (colorOf(uncle) == RED) {
                    setColor(x.parent, BLACK);
                    setColor(uncle, BLACK);
                    setColor(x.parent.parent, RED);
                    x = x.parent.parent;
                } else {
                    if (x == x.parent.right) {
                        leftRightRotate(x.parent.parent);
                        setColor(x, BLACK);
                        setColor(x.right, RED);
                    } else {
                        setColor(x.parent, BLACK);
                        setColor(x.parent.parent, RED);
                        rightRotate(x.parent.parent);
                    }
                }
            } else {
                RBNode<K, V> uncle = x.parent.parent.left;
                if (colorOf(uncle) == RED) {
                    setColor(x.parent, BLACK);
                    setColor(uncle, BLACK);
                    setColor(x.parent.parent, RED);
                    x = x.parent.parent;
                } else {
                    if (x == x.parent.left) {
                        rightLeftRotate(x.parent.parent);
                        setColor(x, BLACK);
                        setColor(x.left, RED);
                    } else {
                        setColor(x.parent, BLACK);
                        setColor(x.parent.parent, RED);
                        leftRotate(x.parent.parent);
                    }
                }
            }
        }
        setColor(root, BLACK);
    }

    private RBNode<K, V> predecessor(RBNode<K, V> node) {
        if (node == null) {
            return null;
        } else if (node.left != null) {
            RBNode<K, V> p = node.left;
            while (p.right != null) {
                p = p.right;
            }
            return p;
        } else {
            RBNode<K, V> p = node.parent;
            RBNode<K, V> ch = node;
            while (p != null && ch == p.left) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    private RBNode<K, V> successor(RBNode<K, V> node) {
        if (node == null) {
            return null;
        } else if (node.right != null) {
            RBNode<K, V> p = node.right;
            while (p.left != null) {
                p = p.left;
            }
            return p;
        } else {
            RBNode<K, V> p = node.parent;
            RBNode<K, V> ch = node;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    public void remove(K key) {
        RBNode<K, V> node = getNode(key);
        if (node == null) return;
        deleteNode(node);
        return;
    }

    private void deleteNode(RBNode<K, V> node) {
        if (node.left != null && node.right != null) {
//            RBNode<K, V> predecessor = predecessor(node);
//            node.value = predecessor.value;
//            node = predecessor;
            RBNode<K, V> successor = successor(node);
            node.key = successor.key;
            node.value = successor.value;
            node = successor;
        }
        RBNode<K, V> replacement = node.left != null ? node.left : node.right;
        if (replacement != null) {
            replacement.parent = node.parent;
            if (node.parent == null) {
                root = replacement;
            } else if (node == node.parent.left) {
                node.parent.left = replacement;
            } else {
                node.parent.right = replacement;
            }
            node.left = node.right = node.parent = null;
            if (colorOf(node) == BLACK) {
                fixAfterRemove(replacement);
            }
        } else if (node.parent == null) {
            root = null;
        } else {
            if (colorOf(node) == BLACK) {
                fixAfterRemove(node);
            }
            if (node.parent != null) {
                if (node == node.parent.left) {
                    node.parent.left = null;
                } else {
                    node.parent.right = null;
                }
            }
        }
    }

    private void fixAfterRemove(RBNode<K, V> x) {
        while (x != root && colorOf(x) == BLACK) {
            if (x == x.parent.left) {
                RBNode<K, V> s = x.parent.right;
                if (colorOf(s) == RED) {
                    leftRotate(x.parent);
                    setColor(s, BLACK);
                    setColor(x.parent, RED);
                    s = x.parent.right;
                }
                if (colorOf(s.left) == BLACK && colorOf(s.right) == BLACK) {
                    setColor(s, RED);
                    x = x.parent;
                } else {
                    if (colorOf(s.right) == BLACK) {
                        RBNode<K, V> p = s.parent;
                        rightLeftRotate(p);
                        setColor(s, BLACK);
                        setColor(s.parent, p.color);
                        setColor(p, BLACK);
                        x = root;
                    } else {
                        RBNode<K, V> p = s.parent;
                        leftRotate(p);
                        setColor(s, colorOf(p));
                        setColor(p, BLACK);
                        setColor(s.right, BLACK);
                        x = root;  // 终止循环
                    }
                }
            } else {
                RBNode<K, V> s = x.parent.left;
                if (colorOf(s) == RED) {
                    setColor(s, BLACK);
                    setColor(x.parent, RED);
                    rightRotate(x.parent);
                    s = x.parent.left;
                }
                if (colorOf(s.left) == BLACK && colorOf(s.right) == BLACK) {
                    setColor(s, RED);
                    x = x.parent;
                } else {
                    if (colorOf(s.left) == BLACK) {
                        RBNode<K, V> p = s.parent;
                        leftRightRotate(p);
                        setColor(s, BLACK);
                        setColor(s.parent, p.color);
                        setColor(p, BLACK);
                        x = root;
                    } else {
                        RBNode<K, V> p = s.parent;
                        rightRotate(p);
                        setColor(s, colorOf(p));
                        setColor(p, BLACK);
                        setColor(s.left, BLACK);
                        x = root;
                    }
                }
            }
        }
        setColor(x, BLACK);
    }

    private RBNode<K, V> getNode(K key) {
        RBNode<K, V> node = this.root;
        while (node != null) {
            int cmp = key.compareTo(node.key);
            if (cmp < 0) {
                node = node.left;
            } else if (cmp > 0) {
                node = node.right;
            } else {
                return node;
            }
        }
        return null;
    }

    private void setColor(RBNode<K, V> node, boolean color) {
        if (node != null) {
            node.color = color;
        }
    }

    private void leftRotate(RBNode<K, V> g) {
        if (g != null) {
            RBNode<K, V> p = g.right;
            g.right = p.left;
            if (p.left != null) {
                p.left.parent = g;
            }
            p.parent = g.parent;
            if (g.parent == null) {
                this.root = p;
            } else if (g.parent.left == g) {
                g.parent.left = p;
            } else {
                g.parent.right = p;
            }
            p.left = g;
            g.parent = p;
        }
    }

    private void rightRotate(RBNode<K, V> g) {
        if (g != null) {
            RBNode<K, V> p = g.left;
            g.left = p.right;
            if (p.right != null) {
                p.right.parent = g;
            }
            p.parent = g.parent;
            if (g.parent == null) {
                this.root = p;
            } else if (g.parent.left == g) {
                g.parent.left = p;
            } else {
                g.parent.right = p;
            }
            p.right = g;
            g.parent = p;
        }
    }

    private void leftRightRotate(RBNode<K, V> g) {
        leftRotate(g.left);
        rightRotate(g);
    }

    private void rightLeftRotate(RBNode<K, V> g) {
        rightRotate(g.right);
        leftRotate(g);
    }
}

class RBNode<K extends Comparable<K>, V> {
    public RBNode<K, V> parent;
    public RBNode<K, V> left;
    public RBNode<K, V> right;
    public boolean color;
    public K key;
    public V value;

    public RBNode(RBNode<K, V> parent, K key, V value) {
        this.parent = parent;
        this.key = key;
        this.value = value;
    }
}

